"""
Skeleton Normalization and Rendering Module

This module handles:
1. Pelvis-centered alignment
2. Torso vertical correction
3. Scale normalization
4. Skeleton line rendering
"""

import cv2
import numpy as np
from typing import Tuple, List
import logging

from .pose_detector import LandmarkIndices

logger = logging.getLogger(__name__)


class SkeletonNormalizer:
    """
    Normalizes skeleton pose for consistent comparison.

    The normalization process:
    1. Center at pelvis (midpoint of hips)
    2. Rotate torso to vertical
    3. Scale to fill canvas proportionally
    """

    # Define skeleton connections for rendering
    SKELETON_CONNECTIONS = [
        # Torso
        (LandmarkIndices.LEFT_SHOULDER, LandmarkIndices.RIGHT_SHOULDER),
        (LandmarkIndices.LEFT_SHOULDER, LandmarkIndices.LEFT_HIP),
        (LandmarkIndices.RIGHT_SHOULDER, LandmarkIndices.RIGHT_HIP),
        (LandmarkIndices.LEFT_HIP, LandmarkIndices.RIGHT_HIP),

        # Left arm
        (LandmarkIndices.LEFT_SHOULDER, LandmarkIndices.LEFT_ELBOW),
        (LandmarkIndices.LEFT_ELBOW, LandmarkIndices.LEFT_WRIST),

        # Right arm
        (LandmarkIndices.RIGHT_SHOULDER, LandmarkIndices.RIGHT_ELBOW),
        (LandmarkIndices.RIGHT_ELBOW, LandmarkIndices.RIGHT_WRIST),

        # Left leg
        (LandmarkIndices.LEFT_HIP, LandmarkIndices.LEFT_KNEE),
        (LandmarkIndices.LEFT_KNEE, LandmarkIndices.LEFT_ANKLE),
        (LandmarkIndices.LEFT_ANKLE, LandmarkIndices.LEFT_HEEL),
        (LandmarkIndices.LEFT_HEEL, LandmarkIndices.LEFT_FOOT_INDEX),

        # Right leg
        (LandmarkIndices.RIGHT_HIP, LandmarkIndices.RIGHT_KNEE),
        (LandmarkIndices.RIGHT_KNEE, LandmarkIndices.RIGHT_ANKLE),
        (LandmarkIndices.RIGHT_ANKLE, LandmarkIndices.RIGHT_HEEL),
        (LandmarkIndices.RIGHT_HEEL, LandmarkIndices.RIGHT_FOOT_INDEX),

        # Neck
        (LandmarkIndices.NOSE, LandmarkIndices.LEFT_SHOULDER),
        (LandmarkIndices.NOSE, LandmarkIndices.RIGHT_SHOULDER),
    ]

    def __init__(self,
                 canvas_size: Tuple[int, int] = (224, 224),
                 scale_fill: float = 0.85,
                 line_thickness: int = 2,
                 line_color: Tuple[int, int, int] = (255, 255, 255),
                 background_color: Tuple[int, int, int] = (0, 0, 0)):
        """
        Initialize the skeleton normalizer.

        Args:
            canvas_size: Output image size (width, height)
            scale_fill: Proportion of canvas to fill with skeleton (0-1)
            line_thickness: Thickness of skeleton lines
            line_color: RGB color for skeleton lines
            background_color: RGB color for background
        """
        self.canvas_size = canvas_size
        self.scale_fill = scale_fill
        self.line_thickness = line_thickness
        self.line_color = line_color
        self.background_color = background_color

    def normalize_landmarks(self, landmarks: np.ndarray) -> np.ndarray:
        """
        Normalize 3D landmarks: center at pelvis, rotate torso vertical, scale.

        Args:
            landmarks: Array of shape (33, 4) [x, y, z, visibility]

        Returns:
            Normalized landmarks of shape (33, 3) [x, y, z]
        """
        # Extract x, y, z coordinates (ignore visibility)
        coords = landmarks[:, :3].copy()

        # Step 1: Center at pelvis (midpoint of left and right hips)
        left_hip = coords[LandmarkIndices.LEFT_HIP]
        right_hip = coords[LandmarkIndices.RIGHT_HIP]
        pelvis = (left_hip + right_hip) / 2.0

        coords -= pelvis

        # Step 2: Rotate torso to vertical
        # Calculate torso vector (from pelvis to shoulders midpoint)
        left_shoulder = coords[LandmarkIndices.LEFT_SHOULDER]
        right_shoulder = coords[LandmarkIndices.RIGHT_SHOULDER]
        shoulders = (left_shoulder + right_shoulder) / 2.0

        torso_vector = shoulders - np.zeros(3)  # pelvis is now at origin

        # Project to 2D (x-y plane) and calculate rotation angle
        torso_2d = torso_vector[:2]
        vertical_2d = np.array([0, 1])  # Target: pointing upward

        # Calculate angle between torso and vertical
        angle = np.arctan2(torso_2d[0], torso_2d[1])

        # Rotation matrix for z-axis rotation (around vertical axis)
        cos_a, sin_a = np.cos(-angle), np.sin(-angle)
        rotation_matrix = np.array([
            [cos_a, -sin_a, 0],
            [sin_a, cos_a, 0],
            [0, 0, 1]
        ])

        coords = coords @ rotation_matrix.T

        # Step 3: Scale normalization
        # Calculate bounding box
        min_coords = coords.min(axis=0)
        max_coords = coords.max(axis=0)
        extent = max_coords - min_coords

        # Scale factor to fit within canvas
        max_extent = extent[:2].max()  # Use x-y extent
        if max_extent > 0:
            scale = self.scale_fill / max_extent
            coords *= scale

        return coords

    def render_skeleton(self,
                       landmarks: np.ndarray,
                       normalized: bool = True) -> np.ndarray:
        """
        Render skeleton as line drawing on canvas.

        Args:
            landmarks: Landmarks array (33, 3 or 4)
            normalized: Whether landmarks are already normalized

        Returns:
            RGB image array of shape (H, W, 3)
        """
        if not normalized:
            if landmarks.shape[1] == 4:
                landmarks = self.normalize_landmarks(landmarks)
            else:
                # Already 3D coordinates
                pass

        # Create canvas
        canvas = np.full(
            (self.canvas_size[1], self.canvas_size[0], 3),
            self.background_color,
            dtype=np.uint8
        )

        # Convert normalized coordinates to pixel coordinates
        # Normalized coords are roughly in range [-0.5, 0.5]
        # Map to canvas: center at (W/2, H/2)
        width, height = self.canvas_size
        pixel_coords = landmarks[:, :2].copy()
        pixel_coords[:, 0] = (pixel_coords[:, 0] + 0.5) * width
        pixel_coords[:, 1] = (0.5 - pixel_coords[:, 1]) * height  # Flip y-axis

        pixel_coords = pixel_coords.astype(np.int32)

        # Draw skeleton connections
        for start_idx, end_idx in self.SKELETON_CONNECTIONS:
            start_pt = tuple(pixel_coords[start_idx])
            end_pt = tuple(pixel_coords[end_idx])

            # Check if points are within canvas bounds
            if (0 <= start_pt[0] < width and 0 <= start_pt[1] < height and
                0 <= end_pt[0] < width and 0 <= end_pt[1] < height):
                cv2.line(canvas, start_pt, end_pt,
                        self.line_color, self.line_thickness)

        # Optionally draw joints as circles
        for pt in pixel_coords:
            if 0 <= pt[0] < width and 0 <= pt[1] < height:
                cv2.circle(canvas, tuple(pt), 2, self.line_color, -1)

        return canvas

    def process_sequence(self,
                        landmarks_sequence: List[np.ndarray],
                        render: bool = True) -> Tuple[List[np.ndarray], List[np.ndarray]]:
        """
        Process a sequence of landmarks.

        Args:
            landmarks_sequence: List of landmark arrays
            render: Whether to render skeleton images

        Returns:
            Tuple of (normalized_sequence, rendered_images)
            - normalized_sequence: List of normalized landmarks
            - rendered_images: List of rendered skeleton images (if render=True)
        """
        normalized_sequence = []
        rendered_images = [] if render else None

        for landmarks in landmarks_sequence:
            if landmarks is None:
                normalized_sequence.append(None)
                if render:
                    rendered_images.append(None)
                continue

            # Normalize
            normalized = self.normalize_landmarks(landmarks)
            normalized_sequence.append(normalized)

            # Render
            if render:
                image = self.render_skeleton(normalized, normalized=True)
                rendered_images.append(image)

        logger.info(f"Processed {len(normalized_sequence)} frames")

        return normalized_sequence, rendered_images


def calculate_joint_angles(landmarks: np.ndarray) -> dict:
    """
    Calculate key joint angles from landmarks.

    Args:
        landmarks: Normalized landmarks (33, 3)

    Returns:
        Dictionary of joint angles in degrees
    """
    def angle_between_vectors(v1: np.ndarray, v2: np.ndarray) -> float:
        """Calculate angle between two vectors in degrees."""
        cos_angle = np.dot(v1, v2) / (np.linalg.norm(v1) * np.linalg.norm(v2) + 1e-8)
        cos_angle = np.clip(cos_angle, -1.0, 1.0)
        return np.degrees(np.arccos(cos_angle))

    angles = {}

    # Left elbow angle
    left_shoulder = landmarks[LandmarkIndices.LEFT_SHOULDER]
    left_elbow = landmarks[LandmarkIndices.LEFT_ELBOW]
    left_wrist = landmarks[LandmarkIndices.LEFT_WRIST]

    v1 = left_shoulder - left_elbow
    v2 = left_wrist - left_elbow
    angles['left_elbow'] = angle_between_vectors(v1, v2)

    # Right elbow angle
    right_shoulder = landmarks[LandmarkIndices.RIGHT_SHOULDER]
    right_elbow = landmarks[LandmarkIndices.RIGHT_ELBOW]
    right_wrist = landmarks[LandmarkIndices.RIGHT_WRIST]

    v1 = right_shoulder - right_elbow
    v2 = right_wrist - right_elbow
    angles['right_elbow'] = angle_between_vectors(v1, v2)

    # Left knee angle
    left_hip = landmarks[LandmarkIndices.LEFT_HIP]
    left_knee = landmarks[LandmarkIndices.LEFT_KNEE]
    left_ankle = landmarks[LandmarkIndices.LEFT_ANKLE]

    v1 = left_hip - left_knee
    v2 = left_ankle - left_knee
    angles['left_knee'] = angle_between_vectors(v1, v2)

    # Right knee angle
    right_hip = landmarks[LandmarkIndices.RIGHT_HIP]
    right_knee = landmarks[LandmarkIndices.RIGHT_KNEE]
    right_ankle = landmarks[LandmarkIndices.RIGHT_ANKLE]

    v1 = right_hip - right_knee
    v2 = right_ankle - right_knee
    angles['right_knee'] = angle_between_vectors(v1, v2)

    # Hip angles (thigh relative to torso)
    pelvis = (left_hip + right_hip) / 2
    left_shoulder_pos = landmarks[LandmarkIndices.LEFT_SHOULDER]
    right_shoulder_pos = landmarks[LandmarkIndices.RIGHT_SHOULDER]
    shoulders = (left_shoulder_pos + right_shoulder_pos) / 2

    torso_vector = shoulders - pelvis

    # Left hip angle
    left_thigh = left_knee - left_hip
    angles['left_hip'] = angle_between_vectors(torso_vector, left_thigh)

    # Right hip angle
    right_thigh = right_knee - right_hip
    angles['right_hip'] = angle_between_vectors(torso_vector, right_thigh)

    return angles
